package tetris.agent;

import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.OpenOption;
import java.util.ArrayList;
import java.util.Collections;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.commons.math3.distribution.MultivariateNormalDistribution;

import tetris.simulator.State;


////
// Implements the Cross Entropy Method for the Tetris problem.
/////
public class AgentEvolver
{
		// Number of agents in the elite set
        int eliteSetSize;
        // A set of samples
        public ArrayList<Agent> agents;
        // The current covariance matrix of the samples.
        public double[][] covariance;
        // The mean of the samples
        public double[] mean;
        // The best sample ever seen
        public double[] best;
        // Apache Math normal distribution for sampling
        public MultivariateNormalDistribution distribution;
        
        
        // These scores are stored for plotting later.
        public ArrayList<Double> bestScores;
        double bestScore = 0;
        
        public AgentEvolver(int numAgents, int numElites, double[] initialMean, double[][] initialCovariance)
        {
                eliteSetSize = numElites;
                initializeAgents(numAgents);
                mean = new double[initialMean.length];
                best = new double[initialMean.length];
                bestScores = new ArrayList<Double>();
                covariance = new double[initialMean.length][initialMean.length];
				Matrix.copy(best, initialMean);
                Matrix.copy(mean, initialMean);
                Matrix. copy(covariance, initialCovariance);
                initializeDistribution();
                resampleAgents();
        }
        
       
        ////
        // Creates the initial distribution of samples.
        ////
        public void initializeDistribution()
        {
        	try
        	{
                distribution = new MultivariateNormalDistribution(mean, covariance);
        	}
        	catch(org.apache.commons.math3.linear.SingularMatrixException exception)
        	{
        		Matrix.inflate(covariance, 0.1);
        		distribution = new MultivariateNormalDistribution(mean, covariance);
        	}
        	catch(org.apache.commons.math3.linear.NonPositiveDefiniteMatrixException exception)
        	{
        		Matrix.inflate(covariance, 0.1);
        		distribution = new MultivariateNormalDistribution(mean, covariance);
        	}
        }
        
        
        ////
        // Fits Sigma, Mu for the elite set given the last run.
        ////
        public void fitEliteSetDistribution(double noise)
        {
                int dim = mean.length;
                // Create the mean
                double[] sumX = new double[dim];
                for(int i = 0; i < eliteSetSize; i++)
                {
                        double[] x = agents.get(i).theta;
                        Matrix.add(sumX, x, sumX);
                }
                
                // The mean is now normalized
                Matrix. scale(sumX, 1.0 / eliteSetSize);
                Matrix. copy(mean, sumX);
                
                // Now we create the covariance matrix
                double[][] sumSigma = new double[dim][dim];
                double[] diff = new double[dim];
                double[][] outer = new double[dim][dim];
                for(int i = 0; i < eliteSetSize; i++)
                {
                        // Summing up squared errors from the mean
                        double[] x = agents.get(i).theta;
                        Matrix.sub(diff, x, mean);
                        Matrix.outerProduct(outer, diff, diff);
                        Matrix.add(sumSigma, outer, sumSigma);
                }
                
                // Scale the covariance matrix to be unbiased
                Matrix.scale(sumSigma, 1.0f / (eliteSetSize - 1));
                Matrix.copy(covariance, sumSigma);
                Matrix.inflate(covariance, noise);

        }
        
        ////
        // Writes scores to a file for plotting.
        ////
        public void writeScores(String path)
        {
        	Charset charset = Charset.forName("US-ASCII");
        	BufferedWriter writer = null;
        	try 
        	{
				writer = Files.newBufferedWriter(new File(path).toPath(), charset);
			} 
        	catch (IOException e)
			{
				e.printStackTrace();
				return;
			}
        	
        	try
        	{
	            for(int i = 0; i < bestScores.size(); i++)
	            {
	                   writer.write(bestScores.get(i).toString() + " ");
	            }
	        	writer.close();
        	}
        	catch(IOException e)
        	{
        		e.printStackTrace();
        		return;
        	}
        }
        
        ////
        // Writes agents to a file for plotting
        ////
        public void writeAgents(String path)
        {
        	Charset charset = Charset.forName("US-ASCII");
        	BufferedWriter writer = null;
        	try 
        	{
				writer = Files.newBufferedWriter(new File(path).toPath(), charset);
			} 
        	catch (IOException e)
			{
				e.printStackTrace();
				return;
			}
        	
        	try
        	{
	            for(int i = 0; i < agents.size(); i++)
	            {
	                   Matrix.write(writer, agents.get(i).theta);
	                   writer.write("\n");
	            }
	        	writer.close();
        	}
        	catch(IOException e)
        	{
        		e.printStackTrace();
        		return;
        	}
        	
        }
        
        ////
        // Given the current distribution, resamples all the agents.
        ////
        public void resampleAgents()
        {
                for(int i = 0; i < agents.size(); i++)
                {
                       Matrix.copy(agents.get(i).theta, distribution.sample());
                }
        }
        
        ////
        // Actually runs the Cross Entropy Method
        ////
        public void evolveAgents(State start, int generations)
        {
        		double noiseValue = 100.0f;
        		double noiseDecrease = 1.0f;
                for(int i = 0; i < generations; i++)
                {
                        System.out.printf("Generation %d/%d\n", i, generations);
                        System.out.printf("Scoring agents\n");
                        scoreAgents(start);
                        System.out.printf("Sorting agents\n");
                        sortAgents();
                        System.out.printf("Best agent has score %f worst: %f median: %f\n", agents.get(0).score, agents.get(agents.size() - 1).score, agents.get(agents.size() / 2).score);

                        
                        double currentScore = agents.get(0).score;
                        bestScores.add(currentScore);
                        
                        if(currentScore > bestScore)
                        {
                        	bestScore = currentScore;
                        	Matrix.copy(best, agents.get(0).theta);
                        }
                        
                        for(int j = 0; j < eliteSetSize; j++)
                        {
                                System.out.printf("%f ", agents.get(j).lines);
                        }
                        
                        System.out.printf("\n");
                        
                        System.out.printf("Fitting agents\n");
                        fitEliteSetDistribution(noiseValue);
                        noiseValue-=noiseDecrease;
                        
                        System.out.printf("Mean:\n");
                        for(int j = 0; j < mean.length; j++)
                        {
                                System.out.printf("%s %f (%f)\n", Agent.feature.values()[j].toString(), mean[j], covariance[j][j]);
                        }
                        
                        
                        System.out.printf("Sampling agents\n");
                        //writeAgents("./agents_" + i + ".txt");
                        initializeDistribution();
                        resampleAgents();
                }
                
                //writeScores("./scores_inflation.txt");

        }
        
        ////
        // Prints every agent's parameters to the screen.
        /////
        public void printAgents()
        {
        	System.out.printf("[");
        	for(int i = 0; i < agents.size(); i++)
        	{
        		Matrix.print(agents.get(i).theta);
        		System.out.printf(" %f\n", agents.get(i).score);
        		
        	}
        	System.out.printf("]");
        		
        		
        }
       
        
        ///
        // Sorts all the agents by their scores.
        ////
        public void sortAgents()
        {
                Collections.sort(agents);
        }
        
        
        
        /////
        /// In parallel, runs all of the agents to completion
        // and scores them.
        /////
        public void scoreAgents(final State start)
        {
                
                
                Parallel.For(agents,
                                                new Parallel.Operation<Agent>()
                                                {
                                                        public void perform(Agent param)
                                                        {
                                                        		
                                                                int numRuns = 4;
                                                                double sumScore = 0;
                                                                double sumLines = 0;
                                                                for(int runs = 0; runs < numRuns; runs++)
                                                                {
                                                                        param.forwardSimulate(start, 1000000, 1.0 - 1e-12);
                                                                        sumScore += param.lines;
                                                                        sumLines += param.lines;
                                                                }
                                                                
                                                                param.score = sumScore / numRuns;
                                                                param.lines = sumLines / numRuns;
                                                                
                                                        };
                                                });
                
            
        }
        
        /// Creates the initial set of agents.
        public void initializeAgents(int numAgents)
        {
                agents = new ArrayList<Agent>();
                
                for(int i = 0; i < numAgents; i++)
                {
                        Agent agent = new Agent();
                        agents.add(agent);
                }
        }
        
        public ArrayList<Agent> getAgents()
        {
                return agents;
        }
        

}
